{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyMlDthhM8w5NIUNYffwmHfr",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sngo/llms-practice/blob/main/taskmanagement/TaskManagement.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "GM-U1_ArmVFO"
      },
      "outputs": [],
      "source": [
        "!pip install -q gradio>=4.0.0"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import gradio as gr\n",
        "import openai\n",
        "import json\n",
        "import html\n",
        "from openai import OpenAI\n",
        "import random\n",
        "import datetime\n",
        "from google.colab import userdata"
      ],
      "metadata": {
        "id": "FciiJrKSmfOV"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "api_key = userdata.get('OPENAI_API_KEY')\n",
        "client = OpenAI(api_key=api_key)\n",
        "model = \"gpt-4.1-mini\""
      ],
      "metadata": {
        "id": "7fPLICmznVur"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def abcd_taskTool():\n",
        "  tasks = [\n",
        "         {\"taskId\": \"T001\", \"accNo\": \"1234567890\", \"status\": \"Pending\", \"description\": \"Update account details\"},\n",
        "        {\"taskId\": \"T002\", \"accNo\": \"9876543210\", \"status\": \"In Progress\", \"description\": \"Verify transaction\"}\n",
        "    ]\n",
        "\n",
        "  probability_of_tasks = 0.8\n",
        "  if random.random() < probability_of_tasks:\n",
        "    return tasks\n",
        "  else:\n",
        "    return []"
      ],
      "metadata": {
        "id": "zyXLchykmoW_"
      },
      "execution_count": 35,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def abcd_NotifyTool():\n",
        "    # Mock notification function\n",
        "    return True"
      ],
      "metadata": {
        "id": "h4J8HOOJm78X"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Tool definitions for OpenAI\n",
        "tools = [\n",
        "    {\n",
        "        \"type\": \"function\",\n",
        "        \"function\": {\n",
        "            \"name\": \"abcd_taskTool\",\n",
        "            \"description\": \"Retrieve a list of ABCD tasks.\",\n",
        "            \"parameters\": {}\n",
        "        }\n",
        "    },\n",
        "    {\n",
        "        \"type\": \"function\",\n",
        "        \"function\": {\n",
        "            \"name\": \"abcd_NotifyTool\",\n",
        "            \"description\": \"Notify the support team about tasks.\",\n",
        "            \"parameters\": {}\n",
        "        }\n",
        "    }\n",
        "]"
      ],
      "metadata": {
        "id": "KOtTI6PunNJ7"
      },
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Function to mask first four digits of accNo\n",
        "def mask_accNo(accNo):\n",
        "    return \"****\" + accNo[4:] if len(accNo) >= 4 else accNo"
      ],
      "metadata": {
        "id": "p7wU9CRBnTAS"
      },
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Function to format tasks as an HTML table\n",
        "def format_tasks_as_table(tasks):\n",
        "    if not tasks:\n",
        "        return \"No tasks found.\"\n",
        "\n",
        "    table = \"<table border='1' style='border-collapse: collapse; width: 100%;'>\"\n",
        "    table += \"<tr><th>Task ID</th><th>Account Number</th><th>Status</th><th>Description</th></tr>\"\n",
        "    for task in tasks:\n",
        "        table += \"<tr>\"\n",
        "        table += f\"<td>{html.escape(task.get('taskId', ''))}</td>\"\n",
        "        table += f\"<td>{html.escape(mask_accNo(task.get('accNo', '')))}</td>\"\n",
        "        table += f\"<td>{html.escape(task.get('status', ''))}</td>\"\n",
        "        table += f\"<td>{html.escape(task.get('description', ''))}</td>\"\n",
        "        table += \"</tr>\"\n",
        "    table += \"</table>\"\n",
        "    return table"
      ],
      "metadata": {
        "id": "b4GLrDLyntA7"
      },
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Chat function to handle user input and bot responses with tool calling\n",
        "def chat_function(message, history):\n",
        "    # Append user message to history\n",
        "    print(f\"User message: {message}\")\n",
        "\n",
        "    history.append([message, None])\n",
        "    print(f\"History: {history}\")\n",
        "\n",
        "    # Prepare messages for OpenAI API\n",
        "    messages = [{\"role\": \"system\", \"content\": (\n",
        "        \"You are a helpful assistant. When the user asks about checking ABCD tasks, call the abcd_taskTool function. \"\n",
        "        \"If tasks are returned, ask the user if they want to notify the support team. \"\n",
        "        \"If the user agrees to notify, call the abcd_NotifyTool function. \"\n",
        "        \"If no tasks are found, respond with a cheerful message. \"\n",
        "        \"If the user declines notification, respond creatively.\"\n",
        "    )}]\n",
        "    for user_msg, bot_msg in history[:-1]:\n",
        "        if user_msg:\n",
        "            messages.append({\"role\": \"user\", \"content\": user_msg})\n",
        "        if bot_msg:\n",
        "            messages.append({\"role\": \"assistant\", \"content\": bot_msg})\n",
        "    messages.append({\"role\": \"user\", \"content\": message})\n",
        "\n",
        "    # Handle state for notification response\n",
        "    state = history[-2][1] if len(history) > 1 else None\n",
        "    if state and \"Would you like to notify the support team?\" in state:\n",
        "        if message.lower() in [\"yes\", \"y\", \"yeah\", \"sure\", \"ok\", \"notify\", \"yep\"]:\n",
        "            # Simulate tool call for notification\n",
        "            notify_result = abcd_NotifyTool()\n",
        "            response = \"The support team has been notified successfully! They'll take it from here. 😊\"\n",
        "            history[-1][1] = response\n",
        "            yield history, \"normal\", \"\"\n",
        "            return\n",
        "        else:\n",
        "            response = (\n",
        "                \"Alright, we'll hold off on notifying the team. 🌟 \"\n",
        "                \"How about we explore something fun, like planning your next big adventure?\"\n",
        "            )\n",
        "            history[-1][1] = response\n",
        "            yield history, \"normal\", \"\"\n",
        "            return\n",
        "\n",
        "    try:\n",
        "        # Call OpenAI API with tool calling\n",
        "        response = client.chat.completions.create(\n",
        "            model= model,\n",
        "            messages=messages,\n",
        "            tools=tools,\n",
        "            tool_choice=\"auto\",\n",
        "            max_tokens=300\n",
        "        )\n",
        "\n",
        "        # Process response\n",
        "        response_message = response.choices[0].message\n",
        "        tool_calls = response_message.tool_calls\n",
        "\n",
        "        if tool_calls:\n",
        "            for tool_call in tool_calls:\n",
        "                function_name = tool_call.function.name\n",
        "                if function_name == \"abcd_taskTool\":\n",
        "                    tasks = abcd_taskTool()\n",
        "                    task_count = len(tasks)\n",
        "\n",
        "                    if task_count > 0:\n",
        "                        table = format_tasks_as_table(tasks)\n",
        "                        response = (\n",
        "                            f\"There are {task_count} tasks found:<br>{table}<br>\"\n",
        "                            \"Would you like to notify the support team? (Yes/No)\"\n",
        "                        )\n",
        "                        history[-1][1] = response\n",
        "                        print(f\"History after tool call: {history}\")\n",
        "                        print(f\"History at -1, 1: {history[-1][1]}\")\n",
        "                        yield history, \"waiting_for_notify_response\", \"\"\n",
        "                    else:\n",
        "                        response = (\n",
        "                            \"Hooray! No ABCD tasks to worry about! 🎈 \"\n",
        "                            \"You're free as a bird—any fun plans for the day?\"\n",
        "                        )\n",
        "                        history[-1][1] = response\n",
        "                        yield history, \"normal\", \"\"\n",
        "                elif function_name == \"abcd_NotifyTool\":\n",
        "                    abcd_NotifyTool()\n",
        "                    response = \"The support team has been notified successfully! They'll take it from here. 😊\"\n",
        "                    history[-1][1] = response\n",
        "                    yield history, \"normal\", \"\"\n",
        "        else:\n",
        "            # No tool calls, use the LLM's response\n",
        "            bot_response = response_message.content\n",
        "            history[-1][1] = bot_response\n",
        "            yield history, \"normal\", \"\"\n",
        "\n",
        "    except Exception as e:\n",
        "        history[-1][1] = f\"Error: {str(e)}\"\n",
        "        yield history, \"normal\", \"\""
      ],
      "metadata": {
        "id": "m9MN8vFDnwsA"
      },
      "execution_count": 48,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Custom CSS for table and chat styling\n",
        "custom_css = \"\"\"\n",
        "table {\n",
        "    font-family: Arial, sans-serif;\n",
        "    margin: 10px 0;\n",
        "}\n",
        "th, td {\n",
        "    padding: 8px;\n",
        "    text-align: left;\n",
        "}\n",
        "th {\n",
        "    background-color: #f2f2f2;\n",
        "}\n",
        "tr:nth-child(even) {\n",
        "    background-color: #f9f9f9;\n",
        "}\n",
        "\"\"\""
      ],
      "metadata": {
        "id": "X3Egf6jAn3uP"
      },
      "execution_count": 31,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Create Gradio interface\n",
        "with gr.Blocks(css=custom_css) as demo:\n",
        "    chatbot = gr.Chatbot()\n",
        "    msg = gr.Textbox(placeholder=\"Type your message here...\")\n",
        "    state = gr.State(value=\"normal\")\n",
        "\n",
        "    msg.submit(\n",
        "        chat_function,\n",
        "        inputs=[msg, chatbot],\n",
        "        outputs=[chatbot, state, msg]\n",
        "    )"
      ],
      "metadata": {
        "id": "g9ctQlBOoDvi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "demo.launch(debug=True, share=True)"
      ],
      "metadata": {
        "id": "DG5gjlXloHhj"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "GiWUiQNJxPSq"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}